# 本文件包含所有与数据获取相关的函数
# by yewentao
# v0.1
import tushare as ts
import _params

# 使用TS获取股票列表信息，数据来自sina财经
def collect_stock_list_by_sina():
    stock_list = ts.get_stock_basics()
    stock_list['code'] = stock_list.index
    stock_list = stock_list.reset_index(drop=True)
    stock_list = stock_list[_params.__STOCK_LIST_FIELDS_SINA]
    # 返回必要的字段
    return stock_list

# 使用通联接口获取股票基本信息，可以获取所有股票列表，也可以获取单个股票的基本信息
# 参数说明见通联API文档：getEqu
def collect_stock_list_by_tl(equTypeCD='', secID='', ticker='', listStatusCD='', field=''):
    ts.set_token(_params.__TL_TOKEN)
    eq = ts.Equity()
    stock_list = eq.Equ(equTypeCD, secID, ticker, listStatusCD, field)
    stock_list = stock_list.fillna('')
    # 针对ticker转为数字的情况，强制转换为六位字符串
    stock_list['ticker'] = stock_list['ticker'].apply(lambda x : str(x).zfill(6))
    return stock_list

# 使用通联接口获取某个股票日线行情数据
# 参数说明见通联API文档：getMktEqud
def collect_stock_day_by_tl(secID = '', ticker = '', tradeDate = '', beginDate='', endDate='', field=''):
    ts.set_token(_params.__TL_TOKEN)
    mkt = ts.Market()
    stock_hq = mkt.MktEqud(secID=secID, ticker=ticker, tradeDate=tradeDate, beginDate=beginDate, endDate=endDate, field=field)
    # 针对ticker转为数字的情况，强制转换为六位字符串
    stock_hq = stock_hq.fillna('')
    stock_hq['ticker'] = stock_hq['ticker'].apply(lambda x : str(x).zfill(6))
    return stock_hq

# 获取当前日期可用的最新交易日，如果当日是交易日，则返回当日，如果不是，则返回前一交易日
# 交易所默认为上交所
def collect_newest_tradedate(date, exchange='XSHG'):
    ts.set_token(_params.__TL_TOKEN)
    master = ts.Master()
    df = master.TradeCal(exchangeCD='XSHG',beginDate=date, endDate=date)
    if df.empty:
        print("获取日历失败，请检查参数是否正确:" + str(date))
        exit()
    df = df.iloc[0]
    if df['isOpen'] == 1:
        return date
    else:
        return df['prevTradeDate'].replace('-','')